LFM2 #20797

paulpak58 · 2025-07-11T05:23:45Z

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

Purpose

This PR implements LFM2.

Implementation Details

These changes are relatively lightweight & independent. Some components were modified to generalize the Mamba2 V1 modules to handle for both MambaSpec’s and ShortConvSpec’s, e.g. HybridAttentionMambaModelConfig -> HybridAttentionStaticCacheModelConfig, which will be needed for other V1 helpers (e.g. MambaManager) once the horizon of static cache/fixed-state layers increases. For simplicity, in most cases, I try to use the same V1 modules to handle for both Mamba2 and ShortConv layers.

Test Plan

See test_hybrid.py.

Test Result

See test_hybrid.py. Basic hybrid tests, batching, chunked_prefill passes. Later tests don't pass due to scheduling incompatibilities with hybrids (model-agnostic).

Need transformers v.4.54.0.dev0 (install from source). This version doesn't support token_type_ids which is populated in HfRunner. I did not change this in the PR but if you want to test the hybrid V1 path locally in test_hybrid.py, you’ll have to pop those kwargs from wrapped_inputs in conftest.py before passing them through the HF model.

(Optional) Documentation Update

Updates to supported_models.md and v1_guide.md.

…with Mamba2

…ormers >= 4.54.0.dev0

…usal_conv1d kernel

…n 4.54.0.dev0

github-actions · 2025-07-11T05:23:54Z

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

gemini-code-assist

Summary of Changes

Hello @paulpak58, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the LFM2 model architecture into vLLM, including necessary components for hybrid attention and convolution-based layers. It also incorporates changes to handle CUDA versions, transformers library updates, and integration with the vLLM v1 framework.

Highlights

CMakeLists.txt: Added logic to ignore nvToolsExt for CUDA 12.9 by creating an imported interface library if it's not already a target.
vllm/config.py: Modified get_num_layers_by_block_type to handle hybrid models, specifically addressing an attribute change in transformers library versions >= 4.54.0.dev0. It now checks for layer_types in hf_text_config if layers_block_type is not found in hf_config. Also, the logic for summing the block types has been updated to account for 'full_attention' blocks when the block type is 'attention'.
vllm/model_executor/layers/conv.py: Introduced a new ShortConv custom operator for implementing a short convolution layer, including forward passes for both native and CUDA execution. The CUDA forward pass includes logic for prefill and decode stages, utilizing causal_conv1d_fn and causal_conv1d_update respectively. The class also defines a get_state_shape method.
vllm/model_executor/models/conv_cache.py: Introduced ConvCacheParams dataclass and ConvCacheManager class to manage the convolution state cache. The ConvCacheManager inherits from ConstantSizeCache and provides methods for copying cache, retrieving tensors for the current run, and providing CUDA graph capture inputs.
vllm/model_executor/models/lfm2.py: Added new modules and classes for the LFM2 model architecture, including LFM2MLP, LFM2Attention, LFM2AttentionDecoderLayer, LFM2ShortConvDecoderLayer, LFM2Model, and LFM2ForCausalLM. These components define the layers and overall structure of the LFM2 model, supporting hybrid attention and convolution-based layers. The LFM2ForCausalLM class integrates the LFM2 model with the vLLM framework, including cache management and logits processing.
vllm/model_executor/models/registry.py: Registered LFM2ForCausalLM in the model registry, associating it with the 'lfm2' identifier.
vllm/transformers_utils/configs/ovis.py: Wrapped the AutoConfig.register call for AIMv2Config in a try-except block to prevent errors if AutoConfig is not available.
vllm/utils/init.py: Added conv to the LayerBlockType enum.
vllm/v1/attention/backends/mamba_attn.py: Added get_short_conv_chunk_size function and updated Mamba2AttentionMetadataBuilder to support ShortConvSpec.
vllm/v1/core/single_type_kv_cache_manager.py: Added ShortConvSpec to the dictionary of kv cache specs.
vllm/v1/kv_cache_interface.py: Added ShortConvSpec dataclass to define the specification for the short convolution KV cache.
vllm/v1/worker/gpu_model_runner.py: Imported ShortConv, updated initialize_attn_backend to support ShortConvSpec, and modified _reshape_kv_cache_tensors and get_kv_cache_spec to handle short convolution layers.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces support for the LFM2 model, a new hybrid architecture. The changes are comprehensive, affecting model definition, caching mechanisms, and the core engine logic. The implementation cleverly reuses existing infrastructure for Mamba-like layers to handle the new convolution layers, which is a great approach.

I've provided several suggestions to enhance code clarity, maintainability, and robustness. These include refactoring duplicated code, simplifying complex conditions, and improving exception handling. Overall, this is a solid contribution.

vllm/transformers_utils/configs/ovis.py

vllm/config.py

vllm/model_executor/layers/conv.py

vllm/v1/kv_cache_interface.py

vllm/v1/worker/gpu_model_runner.py

mergify · 2025-07-11T05:56:40Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @paulpak58.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 · 2025-07-15T21:48:30Z

CMakeLists.txt

+if (NOT TARGET CUDA::nvToolsExt)
+    add_library(CUDA::nvToolsExt INTERFACE IMPORTED)
+endif()
+


possibly a cleaner solution than this, but this works.

mergify · 2025-07-16T02:26:06Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @paulpak58.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Paul Pak <paulpak58@gmail.com>

robertgshaw2-redhat · 2025-07-16T18:27:57Z

cc @tdoublep - could you take a look at this?

Signed-off-by: Paul Pak <paulpak58@gmail.com>

tdoublep

Nice work! My general impression is that this is a pretty clean integration. I've left some comments after an initial pass.

tests/models/language/generation/test_hybrid.py

vllm/v1/worker/gpu_model_runner.py

tdoublep · 2025-07-16T19:02:38Z

vllm/model_executor/models/conv_cache.py

+                                self.state_indices_tensor)
+
+
+class ConvCacheManager(ConstantSizeCache):


Is this class needed for anything other than V0 support? If so, I question whether we should include it.

Yeah only v0. Do you have an estimate for full v0 deprecation?

vllm/model_executor/models/config.py

vllm/model_executor/layers/conv.py

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 · 2025-07-16T20:01:16Z

Thanks @tdoublep -- appreciate the fast turnaround. All comments are addressed. I haven't removed conv_cache.py in order for the v0 tests to pass. I'd prefer to keep it until all legacy files are removed (mamba_cache.py, conv_cache.py). Let me know if you're fine with that.

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 · 2025-07-17T16:34:51Z

vllm/v1/attention/backends/mamba_attn.py

@@ -89,11 +89,17 @@ class Mamba2AttentionMetadataBuilder(

    def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,


kv_cache_spec is not really an AttentionSpec here. It's a StaticCacheSpec but for us to resolve mypy issues, we'd have to do something like the following:

M = TypeVar("M") S = TypeVar("S", bound="AttentionSpec") class AttentionMetadataBuilder(Generic[M, S], abc.ABC): class Mamba2AttentionMetadataBuilder( AttentionMetadataBuilder[Mamba2AttentionMetadata, StaticCacheSpec]):

It's not messy, but I just found it quicker to label kv_cache_spec as an AttentionSpec to make mypy happy and then check the type in the init.

mergify · 2025-07-18T11:03:40Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @paulpak58.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 · 2025-07-19T21:57:09Z

Both eager & torch-compile passes internal tests.

paulpak58 added 7 commits July 10, 2025 19:44

[cmake] ignore nvToolsExt for cuda-12.9

3a50223

[model_executor][models] LFM2 architecture

6fd86d2

[configs] use layer_types from huggingface hybrids >= 4.54.0.dev0

aaf7df1

[model_runner][v1] ShortConvSpec for ShortConv layers; compatibility …

6c80caf

…with Mamba2

[configs] need to detect full_attention key in layer_types for transf…

d17c95f

…ormers >= 4.54.0.dev0

[layers][conv] update ShortConv layer to be compatible with triton ca…

1bc8835

…usal_conv1d kernel

[transformers][ovis] tmp: AIMv2Config doesn't need to be registered o…

e550362

…n 4.54.0.dev0

mergify bot added ci/build new-model Requests to new models v1 labels Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

paulpak58 changed the title ~~Lfm2~~ LFM2 Jul 11, 2025

gemini-code-assist bot reviewed Jul 11, 2025

View reviewed changes

mergify bot added the needs-rebase label Jul 11, 2025

paulpak58 added 4 commits July 11, 2025 16:17

[models][lfm2] LFM2->Lfm2 to match config

05af65a

[merge] upstream @ 5bac613

40d81e9

[v1][cache] add support for conv cache shapes

7241660

[v1][config] generalize HybridAttentionMambaModelConfig

3d3be6a

mergify bot removed the needs-rebase label Jul 15, 2025

paulpak58 added 3 commits July 15, 2025 21:22

[merge] upstream main @ 19c8630

b2447dd

[layer][conv] update conv metadata in causal_conv1d

46902dc

[misc] format + cleanup

260e3fe

Signed-off-by: Paul Pak <paulpak58@gmail.com>

paulpak58 commented Jul 15, 2025

View reviewed changes

mergify bot added the needs-rebase label Jul 16, 2025

paulpak58 added 2 commits July 16, 2025 15:29

[layers][conv] fix minor discprencies in decode conv

1dff6e1

Signed-off-by: Paul Pak <paulpak58@gmail.com>

[merge] upstream @ a0f8a79

30621b4

mergify bot removed the needs-rebase label Jul 16, 2025

paulpak58 requested review from robertgshaw2-redhat, njhill, comaniac, alexm-redhat, simon-mo, youkaichao, mgoin, tlrmchlsmth, houseroad, hmellor and LucasWilkinson as code owners July 16, 2025 18:23

[docs] update supported_models + v1 guide

7577e89

Signed-off-by: Paul Pak <paulpak58@gmail.com>

mergify bot added the documentation Improvements or additions to documentation label Jul 16, 2025

[misc] fix pre-commit checks

1ff0c89

Signed-off-by: Paul Pak <paulpak58@gmail.com>

tdoublep reviewed Jul 16, 2025

View reviewed changes

[model_executor] remap mamba V1 utils to static_cache + cleanup

b425c0d

Signed-off-by: Paul Pak <paulpak58@gmail.com>

[misc] minor: fix format

80a2f3a

Signed-off-by: Paul Pak <paulpak58@gmail.com>

heheda12345 mentioned this pull request Jul 17, 2025

[RFC]: Native support for Mamba, SSM, and hybrid transformer models in vLLM V1 #17140

Open

9 tasks

[merge] upstream @ 9fb2d22

5f2c6c8

paulpak58 commented Jul 17, 2025

View reviewed changes

EwoutH mentioned this pull request Jul 17, 2025

[MODEL EVALUATION REQUEST] LiquidAI/LFM2-1.2B EuroEval/EuroEval#1076

Open

4 tasks

mergify bot added the needs-rebase label Jul 18, 2025

[merge] upstream @ ed8cbfe

cbc6ba3

Signed-off-by: Paul Pak <paulpak58@gmail.com>

mergify bot removed the needs-rebase label Jul 18, 2025

paulpak58 added 2 commits July 19, 2025 20:13

[merge] upstream @ 881e3cb

7318322

Signed-off-by: Paul Pak <paulpak58@gmail.com>

[models][lfm2] torch compile support

d8b7170

Signed-off-by: Paul Pak <paulpak58@gmail.com>

		self.state_indices_tensor)


		class ConvCacheManager(ConstantSizeCache):

		@@ -89,11 +89,17 @@ class Mamba2AttentionMetadataBuilder(

		def __init__(self, kv_cache_spec: AttentionSpec, vllm_config: VllmConfig,

Uh oh!

LFM2 #20797

Are you sure you want to change the base?

LFM2 #20797

Conversation

paulpak58 commented Jul 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Essential Elements of an Effective PR Description Checklist

Purpose

Implementation Details

Test Plan

Test Result

(Optional) Documentation Update

Uh oh!

github-actions bot commented Jul 11, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mergify bot commented Jul 11, 2025

Uh oh!

paulpak58 Jul 15, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Jul 16, 2025

Uh oh!

robertgshaw2-redhat commented Jul 16, 2025

Uh oh!

tdoublep left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tdoublep Jul 16, 2025

Choose a reason for hiding this comment

Uh oh!

paulpak58 Jul 16, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

paulpak58 commented Jul 16, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

paulpak58 Jul 17, 2025

Choose a reason for hiding this comment

Uh oh!

mergify bot commented Jul 18, 2025

Uh oh!

paulpak58 commented Jul 19, 2025

Uh oh!

Uh oh!

paulpak58 commented Jul 11, 2025 •

edited

Loading

paulpak58 commented Jul 16, 2025 •

edited

Loading